#!/usr/bin/env python3
"""
process_agent.py
Lightweight eBPF/BCC agent to capture process lifecycle events (exec, fork, exit).
Fixed version: Simplified eBPF program for better compatibility.
"""
from bcc import BPF
from datetime import datetime, timezone
import json
import os
import signal
import sys

OUT_DIR = "/var/log/os_monitor"
os.makedirs(OUT_DIR, exist_ok=True)
OUT_FILE = os.path.join(OUT_DIR, "process.jsonl")

# 简化版的eBPF程序，更好的兼容性
BPF_PROGRAM = r"""
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

struct data_t {
    u32 pid;
    u32 ppid;
    u64 ts_ns;
    u32 event;  // 使用u32而不是int
    char comm[16];
};

BPF_PERF_OUTPUT(events);

int trace_sched_process_exec(struct pt_regs *ctx) {
    struct data_t data = {};
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    u32 ppid = bpf_get_current_pid_tgid();
    
    data.pid = pid;
    // 简化：使用固定值或通过其他方式获取ppid
    data.ppid = ppid >> 32;  // 近似获取ppid
    data.ts_ns = bpf_ktime_get_ns();
    data.event = 1;  // exec
    bpf_get_current_comm(&data.comm, sizeof(data.comm));
    
    events.perf_submit(ctx, &data, sizeof(data));
    return 0;
}

int trace_sched_process_fork(struct pt_regs *ctx) {
    struct data_t data = {};
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    
    data.pid = pid;
    data.ppid = 0;  // 简化处理
    data.ts_ns = bpf_ktime_get_ns();
    data.event = 2;  // fork
    bpf_get_current_comm(&data.comm, sizeof(data.comm));
    
    events.perf_submit(ctx, &data, sizeof(data));
    return 0;
}
"""

try:
    b = BPF(text=BPF_PROGRAM)
    
    # 使用tracepoint，更稳定
    b.attach_tracepoint(tp="sched:sched_process_exec", fn_name="trace_sched_process_exec")
    b.attach_tracepoint(tp="sched:sched_process_fork", fn_name="trace_sched_process_fork")
    
except Exception as e:
    print(f"eBPF initialization failed: {e}")
    print("Trying fallback kprobes...")
    
    # 回退到kprobe
    try:
        b.attach_kprobe(event="__x64_sys_execve", fn_name="trace_sched_process_exec")
    except:
        pass

def write_record(record: dict):
    with open(OUT_FILE, "a") as f:
        f.write(json.dumps(record, default=str) + "\n")

def print_event(cpu, data, size):
    try:
        event = b["events"].event(data)
        r = {
            "pid": int(event.pid),
            "ppid": int(event.ppid),
            "comm": event.comm.decode('utf-8', 'replace').strip('\x00'),
            "ts_ns": int(event.ts_ns),
            "event": int(event.event),
            "ts": datetime.now(timezone.utc).isoformat()
        }
        write_record(r)
    except Exception as e:
        print(f"Error processing event: {e}")

b["events"].open_perf_buffer(print_event)
print("[+] process_agent running, writing to {}".format(OUT_FILE))

def handle_sig(signum, frame):
    print("shutting down")
    sys.exit(0)

signal.signal(signal.SIGINT, handle_sig)
signal.signal(signal.SIGTERM, handle_sig)

try:
    while True:
        b.perf_buffer_poll()
except KeyboardInterrupt:
    print("\nShutting down...")